#include "../../Helix.h"
#include "../../HObject.h"
#include "../../Log.h"
#include "Image.h"
#include "PNGImage.h"
#include <string>
#include <iostream>
#include <fstream>
#include <zlib/zlib.h>

using namespace std;

namespace helix
{
	//All PNG images must start with these 8 bytes
	const unsigned char PNGImage::PNG_HEADER[8] = {137,80,78,71,13,10,26,10};
	const unsigned char PNGImage::PNG_IHDR_CHUNK_STRING[4] = {'I','H','D','R'};
	const unsigned char PNGImage::PNG_IDAT_CHUNK_STRING[4] = {'I','D','A','T'};
	
	PNGImage::PNGImage(const char *fileName) : Image(fileName)
	{
		string fName = fileName;
		loaded = loadPNG(&fName);
		depth = 1;		//PNG images are only 2Dimensional

		if(loaded)
		{
			imageType = "PNG";
		}
	}

	PNGImage::~PNGImage()
	{
		if(this->data!=NULL && loaded)
			delete this->data;
	}

	bool PNGImage::loadPNG(string *fileName)
	{
		unsigned char pngH[8];		//the file header

		//Open the file for binary reading
		ifstream fileIn(fileName->data(),ifstream::binary);

		//Make sure the file is open
		if(!fileIn.is_open())
		{
			Log().Error("File: %s did not open correctly. ",fileName->c_str());
			return false;
		}
		
		//read in the PNGHeader (8 bytes)
		fileIn.seekg(0,ios_base::beg);
		fileIn.read((char*)&pngH, 8);

		if(!isPNG(pngH))
		{
			Log().Error("File: %s is not a valid PNG.",fileName->c_str());
			fileIn.close();
			return false;
		}

		//The first chunk is the IHDR, so read it in
		PNGChunk headerChunk;
		if(!readHeaderInfo(&headerChunk,&fileIn))
		{
			Log().Error("File: %s has a header chunk that is not supported.",fileName->c_str());
			fileIn.close();
			return false;
		}

		//Read the next chunk.  We will search until we find the first IDAT chunk
		PNGChunk* dataChunk = new PNGChunk;
		if(!readChunk(dataChunk,&fileIn))
		{
			if(dataChunk->length > 0)
				delete[] dataChunk->data;
			delete dataChunk;
			Log().Error("There was a problem reading the data for file: %s ",fileName->c_str());
			fileIn.close();
			return false;
		}

		bool readSucceeded=true;
		while(getChunkType(dataChunk)!=IDAT_CHUNK && readSucceeded)
		{
			if(dataChunk->length > 0)
				delete[] dataChunk->data;
			delete dataChunk;
			dataChunk = new PNGChunk;
			readSucceeded = readChunk(dataChunk,&fileIn);
		}

		if(!readSucceeded)
		{
			if(dataChunk->length > 0)
				delete[] dataChunk->data;
			delete dataChunk;
			Log().Error("There was a problem reading the data for file: %s ",fileName->c_str());
			fileIn.close();
			return false;
		}

		//At this point, we have the first IDAT chunk
		//uLongf currentLength = 0;
		vector<unsigned char> compressedData;
		//Get the data from the rest of the chunks
		while(getChunkType(dataChunk)==IDAT_CHUNK && readSucceeded)
		{
			/*
			unsigned char *tempArray = new unsigned char[currentLength+dataChunk->length];
			for(unsigned int i=0;i<currentLength;i++)
				tempArray[i]=compressedData[i];
			for(unsigned int i=currentLength;i<currentLength+dataChunk->length;i++)
				tempArray[i]=dataChunk->data[i-currentLength];
			if(currentLength>0)
				delete[] compressedData;
			
			compressedData = tempArray;
			currentLength+=dataChunk->length;
			*/

			for(unsigned int i=0;i<dataChunk->length;i++)
				compressedData.push_back(dataChunk->data[i]);

			if(dataChunk->length > 0)
				delete[] dataChunk->data;
			delete dataChunk;
			dataChunk = new PNGChunk;
			readSucceeded = readChunk(dataChunk,&fileIn);
		}
		if(dataChunk->length > 0)
			delete[] dataChunk->data;
		delete dataChunk;
		if(!readSucceeded || compressedData.size()==0)
		{
			Log().Error("There was a problem reading the data for file: %s ",fileName->c_str());
			fileIn.close();
			return false;
		}

		//At this point, compressedData contains the entire compressed image data (deflated)
		unsigned int unfilteredLength = height * width * bpp/8 + width;
		vector<unsigned char> unfilteredLines(unfilteredLength);
		int zlibRet = uncompressData(&compressedData[0], (unsigned int)compressedData.size(), &unfilteredLines[0], &unfilteredLength);
		
		if(Z_OK != zlibRet)
		{
			Log().Error("There was a problem uncompressing the image data for file: %s ",fileName->c_str());
			fileIn.close();
			return false;
		}

		//At this point, we have the uncompressed filtered scanlines
		//the first byte of each line is the filter type
		int sizeOfData = height * width * bpp/8;
		this->data = new vector<unsigned char>(sizeOfData);
		int bytesPerLine = width* bpp/8;
		for(unsigned int i=0;i<height;i++)
		{
			switch(unfilteredLines[(bytesPerLine+1)*i])
			{
				case 0: noFilter(&unfilteredLines[0],i);
						break;
				case 1: leftFilter(&unfilteredLines[0],i);
						break;
				case 2: upFilter(&unfilteredLines[0],i);
						break;
				case 3: averageFilter(&unfilteredLines[0],i);
						break;
				case 4: paethFilter(&unfilteredLines[0],i);
						break;
				default: //this should not have happened
						delete data;
						data=NULL;
						Log().Error("The filter type is not defined for line: %d ",i);
						fileIn.close();
						return false;
			}
		}

		//Now the data has been read, decompressed, and unfiltered
		//But the image data is stored so that the first scanline is the top
		//of the image.  So we need to flip the data vertically.
		for(int i=0;i<(int)height/2;i++)
		{
			for(int j=0;j<bytesPerLine;j++)
			{
				unsigned char tempByte = (*data)[i*bytesPerLine + j];
				(*data)[i*bytesPerLine + j] = (*data)[(height-i-1)*bytesPerLine + j];
				(*data)[(height-i-1)*bytesPerLine + j] = tempByte;
			}
		}
		//close the file and return
		
		fileIn.close();
		return true;

	}

	bool PNGImage::isPNG(unsigned char* pngH)
	{
		for(int i=0;i<8;i++)
		{
			if(pngH[i] != PNG_HEADER[i])
				return false;
		}

		return true;
	}


	bool PNGImage::readHeaderInfo(PNGChunk *headerChunk,ifstream *fileIn)
	{
		bool readSucceeded = readChunk(headerChunk,fileIn);	//read all the header data into the chunk
		if(headerChunk->length!=13 || !readSucceeded)		//check to make sure the length of the header data is 13 bytes
		{
			Log().Error("IHDR chunk for file is not valid.");
			return false;
		}
		width = getUIntFromByteArray(headerChunk->data,0);		//get the width of the image
		height = getUIntFromByteArray(headerChunk->data,4);		//get the height of the image

		//Check to make sure this is the IHDR chunk, then check bit depth, 
		//color type, compression, filtering, and interlacing
		if(getChunkType(headerChunk)!=IHDR_CHUNK)
		{
			Log().Error("Invalid PNG format, IHDR is not first chunk.");
			return false;
		}
		bpp = (unsigned int)headerChunk->data[8];
		unsigned char colorType = headerChunk->data[9];
		unsigned char compressionType = headerChunk->data[10];
		unsigned char filterType = headerChunk->data[11];
		unsigned char interlaceType = headerChunk->data[12];

		//Make sure this is a true color image (possibly with alpha)
		if(colorType != 2 && colorType != 6)
		{
			Log().Error("Invalid color type, only truecolor is supported.");
			return false;
		}
		//Make sure compression type is zero (this is the only type supported by the PNG specification)
		if(compressionType !=0)
		{
			Log().Error("Invalid compression type.");
			return false;
		}
		//Make sure filter type is zero (this is the only type supported by the PNG specification)
		if(filterType != 0)
		{
			Log().Error("Invalid filter type.");
			return false;
		}
		//Make sure the image is not interlaced (the other type is 1: Adam7 interlacing... this will probably need to be supported)
		if(interlaceType != 0)
		{
			Log().Error("Invalid interlace type, only non-interlaced images are supported.");
			return false;
		}

		//The image is either RGB or RGBA
		int bytesPerChannel = bpp / 8;

		if(colorType == 2)
			channels = 3;
		else 
			channels = 4;

		bpp = bytesPerChannel * channels * 8;

		//unallocate the memory for the data array
		delete[] headerChunk->data;

		return true;
	}

	//Read a chunk from the file into the chunk structure
	//Assumes fileIn's file pointer is pointing to the chunk name
	bool PNGImage::readChunk(PNGChunk *chunk,ifstream *fileIn)
	{
		//Read in the length of the data section of the chunk
		unsigned char length[4];
		fileIn->read((char*)length,4);
		chunk->length = getUIntFromByteArray(length, 0);
		//Read in the chunk type first
		fileIn->read((char*)chunk->type,4);
		//Allocate space for the data contained in the chunk
		if(chunk->length > 0)
			chunk->data = new unsigned char[chunk->length];
		else
			chunk->data = NULL;

		if(chunk->length > 0)
		{
			//Read the data into the chunk
			fileIn->read((char*)chunk->data,chunk->length);
		}

		//Finally, read in the CRC
		fileIn->read((char*)chunk->crc,4);

		if(fileIn->good())
			return true;
		else
			return false;
	}

	
	//Given a byte array and a starting index, parse the next 4 bytes into an unsigned integer.
	unsigned int PNGImage::getUIntFromByteArray(unsigned char *data, unsigned int index)
	{
		unsigned int retVal = ((unsigned int)data[index]<<24) + 
							  ((unsigned int)data[index+1]<<16) + 
							  ((unsigned int)data[index+2]<<8) + 
							  (unsigned int)data[index+3];
		return retVal;
	}

	//Given a chunk, return the type of chunk IHDR, 
	unsigned char PNGImage::getChunkType(PNGChunk *chunk)
	{
		if(isChunkType(chunk,PNG_IHDR_CHUNK_STRING))
			return IHDR_CHUNK;
		if(isChunkType(chunk,PNG_IDAT_CHUNK_STRING))
			return IDAT_CHUNK;
		else
			return UNKNOWN_CHUNK;
	}

	//Return true if all characters in chunk->type[] equal the characters in chunkType[]
	bool PNGImage::isChunkType(PNGChunk *chunk,const unsigned char* chunkType)
	{
		for(int i=0;i<4;i++)
		{
			if(chunk->type[i] != chunkType[i])
				return false;
		}
		return true;
	}


	int PNGImage::uncompressData(unsigned char *source, unsigned int sourceLen, unsigned char *dest, unsigned int *destLen)
	{
		z_stream stream;
		int err;

		stream.next_in = source;
		stream.avail_in = (uInt)sourceLen;
		/* Check for source > 64K on 16-bit machine: */
		if ((uLong)stream.avail_in != sourceLen) return Z_BUF_ERROR;

		stream.next_out = dest;
		stream.avail_out = (uInt)*destLen;
		if ((uLong)stream.avail_out != *destLen) return Z_BUF_ERROR;

		stream.zalloc = (alloc_func)Z_NULL;
		stream.zfree = (free_func)Z_NULL;
		stream.opaque = Z_NULL;

		err = inflateInit(&stream);
		if (err != Z_OK) return err;

		err = inflate(&stream, Z_FINISH);
		if (err != Z_STREAM_END) {
			inflateEnd(&stream);
			if (err == Z_NEED_DICT || (err == Z_BUF_ERROR && stream.avail_in == 0))
				return Z_DATA_ERROR;
			return err;
		}
		*destLen = stream.total_out;

		err = inflateEnd(&stream);
		return err;
	}

	//There is no filtering, the value of the image byte is the source value
	void PNGImage::noFilter(unsigned char *unfilteredLines,unsigned int line)
	{
		int bytes = width*bpp/8;
		for(int i=0;i<bytes;i++)
		{
			(*data)[line*bytes+i]=unfilteredLines[line*(bytes+1)+i+1];
		}
	}

	//There pixel is filtered with the pixel to the left
	void PNGImage::leftFilter(unsigned char *unfilteredLines,unsigned int line)
	{
		int bytespp = bpp/8;
		int bytes = width*bytespp;
		for(int i=0;i<bytes;i++)
		{
			(*data)[line*bytes+i]=unfilteredLines[line*(bytes+1)+i+1] + dataAt(line,i-bytespp);
		}
	}

	//The pixel is filtered with the pixel on the scanline above
	void PNGImage::upFilter(unsigned char *unfilteredLines,unsigned int line)
	{
		int bytespp = bpp/8;
		int bytes = width*bytespp;
		for(int i=0;i<bytes;i++)
		{
			(*data)[line*bytes+i]=unfilteredLines[line*(bytes+1)+i+1] + dataAt(line-1,i);
		}
	}

	//The pixel is filtered with the average of the filter above and to the left
	void PNGImage::averageFilter(unsigned char *unfilteredLines,unsigned int line)
	{
		int bytespp = bpp/8;
		int bytes = width*bytespp;
		for(int i=0;i<bytes;i++)
		{
			int leftPlusUp = (int)dataAt(line,i-bytespp) + (int)dataAt(line-1,i);
			leftPlusUp>>=1;		//divide by 2
			(*data)[line*bytes+i]=unfilteredLines[line*(bytes+1)+i+1] + (unsigned char)(leftPlusUp);
		}
	}

	//The pixel is filtered with the closest pixel out of the ones to the left, top, and top left
	void PNGImage::paethFilter(unsigned char *unfilteredLines,unsigned int line)
	{
		int bytespp = bpp/8;
		int bytes = width*bytespp;
		for(int i=0;i<bytes;i++)
		{
			int a = (int)dataAt(line,i-bytespp);
			int b = (int)dataAt(line-1,i);
			int c = (int)dataAt(line-1,i-bytespp);
			int p = a + b - c;
			int pa = p>a ? p-a : a-p;
			int pb = p>b ? p-b : b-p;
			int pc = p>c ? p-c : c-p;
			int pR;
			if (pa <= pb && pa <= pc) 
				pR = a;
			else if (pb <= pc) 
				pR = b;
			else 
				pR = c;
			
			(*data)[line*bytes+i]=unfilteredLines[line*(bytes+1)+i+1] + (unsigned char)pR;
		}
	}

	unsigned char PNGImage::dataAt(int line,int position)
	{
		if(line < 0)
			return 0;
		if(position < 0)
			return 0;

		int bytespp = bpp/8;
		int bytes = width*bytespp;
		return (*data)[line * bytes + position];
	}
}



